package com.wave.wavesystem.ai.controller;

import com.wave.common.model.Result;
import com.wave.common.util.Func;
import com.wave.wavesystem.ai.config.SSEConfig;
import com.wave.wavesystem.ai.enums.SSEMsgType;
import com.wave.wavesystem.ai.pojo.dto.AiChatDTO;
import com.wave.wavesystem.ai.pojo.entity.AiConversation;
import com.wave.wavesystem.ai.service.IAiConversationService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping("/aiChat")
@RequiredArgsConstructor
public class AiChatController {

    private final ChatClient waveChatClient;
    private final IAiConversationService conversationService;

    @GetMapping(path = "/doConnect", produces = {MediaType.TEXT_EVENT_STREAM_VALUE})
    public SseEmitter doConnect(@RequestParam("conversationId") String conversationId) {
        return SSEConfig.connect(conversationId);
    }

    @PostMapping("/normalChat")
    public Result<String> normalChat(@RequestBody AiChatDTO chatDTO) {
        String conversationId;
        if (Func.isEmpty(chatDTO.getConversationId())) {
            conversationId = conversationService.genConversationId();
            AiConversation conversation = new AiConversation();
            conversation.setConversationId(conversationId);
            conversation.setTitle(chatDTO.getQuery());
            conversationService.save(conversation);
            // 开启对话连接
            SSEConfig.connect(conversationId);
        } else {
            conversationId = chatDTO.getConversationId();
        }
        final String tempCid = conversationId;
        Flux<String> content = waveChatClient
                .prompt(chatDTO.getQuery())
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, tempCid))
                .stream().content();
        content.doFirst(() -> {
                    SSEConfig.sendMessage(tempCid, "", SSEMsgType.START);
                })
                .doOnNext(item -> {
                    SSEConfig.sendMessage(tempCid, item, SSEMsgType.ADD);
                    System.out.println(item + "------------------");
                })
                .doOnComplete(() -> {
                    // 流完成后发送结束信号
                    SSEConfig.sendMessage(tempCid, SSEMsgType.DONE.getValue(), SSEMsgType.DONE);
                })
                .subscribe(); // 需要订阅才会执行
        // 返回过去 开启 sse 连接
        return Result.success(conversationId);
    }


}


